package com.vf.cloud.paas.app.agent.common.util;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.LinkedList;
import java.util.List;

import org.dom4j.Document;
import org.dom4j.DocumentException;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import com.alibaba.fastjson.JSONObject;
import com.jfinal.kit.StrKit;
import com.vf.cloud.paas.common.vo.GpuInfo;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class GpuUtil {

	public static List<GpuInfo> getGpuInfos(String installPath){
		List<GpuInfo> listForGpuInfo=new LinkedList<GpuInfo>();
		String xml = sendGPUCmd(installPath);
		GpuInfo gpuInfo;
		if (!StrKit.isBlank(xml)) {
			List<JSONObject> list = parsingXML(xml);
			for (JSONObject obj : list) {
				gpuInfo = new GpuInfo();
				int index=Integer.parseInt(obj.getString("gpu_module_id"));
				if(index>0) {
					index=index-1;
				}
				gpuInfo.setProductName(obj.getString("product_name"));
				gpuInfo.setMemoryTotal(obj.getString("memoryTotal"));
				gpuInfo.setVbiosVersion(obj.getString("vbios_version"));
				gpuInfo.setModuleId(index+"");
				gpuInfo.setMemoryFree(obj.getString("memoryFree"));
				gpuInfo.setMemoryUsed(obj.getString("memoryUsed"));
				gpuInfo.setBoardId(obj.getString("uuid"));
				
				listForGpuInfo.add(gpuInfo);
			}

		}
		return listForGpuInfo;
	}
	
	private static String sendGPUCmd(String softPath) {
		Process process = null;
		try {
			process = Runtime.getRuntime().exec(String.format("%s/graphics/external/nvidia-smi -q -x", softPath));
			process.getOutputStream().close();
			BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
			StringBuffer stringBuffer = new StringBuffer();
			String line = "";
			while (null != (line = reader.readLine())) {
				stringBuffer.append(line + "\n");
			}
			return stringBuffer.toString();
		} catch (IOException e) {
			log.error("getGPU Error", e);
		}
		return null;
	}

	@SuppressWarnings("unchecked")
	private static List<JSONObject> parsingXML(String s) {
		List<JSONObject> list = new LinkedList<JSONObject>();
		InputStream stream = new ByteArrayInputStream(s.getBytes(StandardCharsets.UTF_8));
		try {
			SAXReader reader = new SAXReader();
			reader.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);
			Document document = reader.read(stream);
			Element root = document.getRootElement();
			List<Element> elements = root.elements();
			JSONObject gpuVo;
			for (Element element : elements) {
				if (StrKit.equals("timestamp", element.getQName().getName())) {

				} else if (StrKit.equals("driver_version", element.getQName().getName())) {

				} else if (StrKit.equals("cuda_version", element.getQName().getName())) {

				} else if (StrKit.equals("attached_gpus", element.getQName().getName())) {

				} else if (StrKit.equals("gpu", element.getQName().getName())) {
					List<Element> gpu = element.elements();
					if (gpu != null && gpu.size() > 0) {
						gpuVo = new JSONObject();
						for (Element e : gpu) {
							if (StrKit.equals("product_name", e.getQName().getName())) {
								gpuVo.put("product_name", e.getText());
							} else if (StrKit.equals("uuid", e.getQName().getName())) {
								gpuVo.put("uuid", e.getText());
							} else if (StrKit.equals("board_id", e.getQName().getName())) {
								gpuVo.put("board_id", e.getText());
							} else if (StrKit.equals("vbios_version", e.getQName().getName())) {
								gpuVo.put("vbios_version", e.getText());
							} else if (StrKit.equals("gpu_module_id", e.getQName().getName())) {
								gpuVo.put("gpu_module_id", e.getText());
							} else if (StrKit.equals("fb_memory_usage", e.getQName().getName())) {
								gpuVo.put("memoryTotal", e.element("total").getText());
								gpuVo.put("memoryUsed", e.element("used").getText());
								gpuVo.put("memoryFree", e.element("free").getText());
							} else if (StrKit.equals("utilization", e.getQName().getName())) {
							}
						}
						list.add(gpuVo);
					}
				}
			}
		} catch (DocumentException e) {
			e.printStackTrace();
		} catch (SecurityException e) {
			e.printStackTrace();
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			try {
				stream.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		return list;
	}

}
